package main

import (
	"bufio"
	"encoding/binary"
	"io"
	"sync"
)

// Frame
//  ------------------------
// | size (2B) | proto (2B) |
//  ------------------------
// |        content         |
//  ------------------------

const (
	frame_Size        = 4
	frame_OffsetSize  = 0
	frame_OffsetProto = 2
)

type Frame struct {
	Size    uint16
	Proto   uint16
	Content []byte
}

///////////////////////////////////////////////////////////////////////////////////
//////////////

type Codec struct {
	r *bufio.Reader
	w *bufio.Writer
	h []byte // header
	p *sync.Pool
}

func NewCodec(rw io.ReadWriter, rbs int, wbs int) *Codec {
	return &Codec{
		r: bufio.NewReaderSize(rw, rbs),
		w: bufio.NewWriterSize(rw, wbs),
		h: make([]byte, frame_Size),
		p: &sync.Pool{New: func() interface{} { return new(Frame) }},
	}
}

func (this *Codec) GetFrame() *Frame {
	return this.p.Get().(*Frame)
}

func (this *Codec) PutFrame(f *Frame) {
	f.Content = nil
	this.p.Put(f)
}

func (this *Codec) Read() (*Frame, error) {
	h, err := this.r.Peek(frame_Size)
	if err != nil {
		return nil, err
	}

	frame := this.GetFrame()
	frame.Size = binary.BigEndian.Uint16(h[frame_OffsetSize:])
	frame.Proto = binary.BigEndian.Uint16(h[frame_OffsetProto:])
	total := int(frame.Size) + frame_Size
	if frame.Size > 0 {
		frame.Content = this.realloc(frame.Content, int(frame.Size))
		b, err := this.r.Peek(total)
		if err != nil {
			return nil, err
		}
		copy(frame.Content, b[frame_Size:])
	}
	this.r.Discard(total)
	return frame, nil
}

func (this *Codec) Write(frame *Frame) error {
	frame.Size = uint16(len(frame.Content))
	binary.BigEndian.PutUint16(this.h[frame_OffsetSize:], frame.Size)
	binary.BigEndian.PutUint16(this.h[frame_OffsetProto:], frame.Proto)

	if _, err := this.w.Write(this.h); err != nil {
		return err
	}

	if frame.Size > 0 {
		if _, err := this.w.Write(frame.Content); err != nil {
			return err
		}
	}
	return this.w.Flush()
}

func (this *Codec) realloc(b []byte, length int) []byte {
	if cap(b) >= length {
		return b[:length]
	}
	alloc := 64
	for alloc < length {
		alloc *= 2
	}
	b = make([]byte, alloc)
	return b[:length]
}
